package com.dmp.storm.base;

import backtype.storm.generated.GlobalStreamId;
import backtype.storm.grouping.CustomStreamGrouping;
import backtype.storm.task.WorkerTopologyContext;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

public abstract class BaseGrouping implements CustomStreamGrouping {
    private static final long serialVersionUID = 1L;
    private Map<String, List<Integer>> groupingMap = new HashMap<String, List<Integer>>();
    private String[] groups;

    public String[] getGroups() {
        return groups;
    }

    public void setGroups(String[] groups) {
        this.groups = groups;
    }

    @Override
    public List<Integer> chooseTasks(int arg0, List<Object> arg1) {
        List<Integer> chooseTask = chooseTasks(groupingMap);
        return chooseTask;
    }

    @Override
    public void prepare(WorkerTopologyContext ctx, GlobalStreamId arg1,
                        List<Integer> arg2) {
        for (String taskName : groups) {
            groupingMap.put(taskName, ctx.getComponentTasks(taskName));
        }
    }

    public abstract List<Integer> chooseTasks(
            Map<String, List<Integer>> groupingMap);

}
